"""
Classification demo for MNIST using the NeuralNetwork class.

To run this demo, use:
    python -m demos.mnist_digits.demo-mnist-digits
from the project root directory.
"""
import numpy as np
import gzip
import os

# Import the neural network components from the new package
from neural_network import Layer
from neural_network import ReLUActivation
from neural_network import SigmoidActivation
from neural_network import NeuralNetwork
from neural_network import LossSoftmaxCCE
# Local import for the draw_digits function
from demos.mnist_digits.draw_digits import draw_digit

# Settings
hidden_sizes = [4, 4]
# hidden_sizes = [128]*3
activation_hidden = ReLUActivation
epochs = 1000
learning_rate = 0.01
batch_size = 64

# Load MNIST from local CSV files
# from the zip file here: https://www.kaggle.com/datasets/oddrationale/mnist-in-csv/data
print("Loading MNIST dataset from local CSV files...")
# Use package-relative paths to the data files
data_dir = os.path.dirname(__file__)
train_path = os.path.join(data_dir, 'mnist_train.csv.gz')
test_path = os.path.join(data_dir, 'mnist_test.csv.gz')
train_data = np.loadtxt(gzip.open(train_path, 'rt'), delimiter=',', skiprows=1)
test_data = np.loadtxt(gzip.open(test_path, 'rt'), delimiter=',', skiprows=1)
print(f"Training samples: {train_data.shape[0]}, Test samples: {test_data.shape[0]}")

# Split features and labels into training data and test data
y_raw_train = train_data[:, 0].astype(int)
X_train = train_data[:, 1:].T / 255.0
y_raw_test = test_data[:, 0].astype(int)
X_test = test_data[:, 1:].T / 255.0

num_classes = 10
# One-hot encode labels
y_train = np.eye(num_classes)[y_raw_train].T
y_test = np.eye(num_classes)[y_raw_test].T

# Build the network
input_size = X_train.shape[0]
network = NeuralNetwork()
prev_size = input_size
for size in hidden_sizes:
    network.add_layer(Layer(prev_size, size))
    network.add_layer(activation_hidden())
    prev_size = size

# Output layer (no activation, Softmax included in loss)
network.add_layer(Layer(prev_size, num_classes))

# Loss function
loss_obj = LossSoftmaxCCE()

# Now we can use the enhanced train method of the NeuralNetwork class
# which supports validation metrics (accuracy or RMSE)
print("Training the neural network...")
network.train(
    X_train, y_train,
    loss_obj,
    epochs,
    learning_rate,
    batch_size,
    shuffle=True,
    verbose=True,
    X_val=X_test,
    y_val=y_test,
    metric="accuracy"
)

# Get final predictions for later use
y_pred_test = network.forward_pass(X_test)

# Now training is done, we'll show off by trying some more digits.
# Interactive digit drawing & classification after training
print("Training complete. Enter interactive digit drawing. Close the drawing window without saving to quit.")
while True:
    digit_vector = draw_digit()
    if digit_vector is None:
        print("No digit received. Exiting interactive mode.")
        break
    # Normalize drawn pixel values (already 0–1)
    x = digit_vector.reshape(-1, 1)
    logits = network.forward_pass(x)
    exp_values = np.exp(logits - np.max(logits, axis=0, keepdims=True))
    probs = exp_values / np.sum(exp_values, axis=0, keepdims=True)
    # Display logits and probabilities in a table
    print(f"{'Class':^5} | {'Logit':^12} | {'Prob (%)':^10}")
    print('-'*31)
    for idx, (l, p) in enumerate(zip(logits.flatten(), probs.flatten())):
        print(f"{idx:^5} | {l:>12.4f} | {p*100:>9.2f}%")